package com.md.config.database.slave;

import com.md.base.exception.ServerException;
import com.md.config.database.slave.balancer.ISlaveSourceLoadBalancer;
import com.md.config.database.slave.balancer.WeightSlaveSourceLoadBalancer;
import com.md.util.CollectionUtil;
import com.md.util.ObjectUtil;
import com.md.value.enums.ErrorEnum;
import com.md.value.server.DataSourceContextHolder;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;
import java.util.Optional;

public class DataSourceRouting extends AbstractRoutingDataSource {
  private final SlaveDataSourceConfig slaveDataSourceConfig;
  private final ISlaveSourceLoadBalancer slaveSourceLoadBalancer;

  public DataSourceRouting(final SlaveDataSourceConfig slaveDataSourceConfig,
                           final ISlaveSourceLoadBalancer slaveSourceLoadBalancer) {
    if (ObjectUtil.isEmpty(slaveDataSourceConfig)) {
      throw ServerException.build(ErrorEnum.BEAN_REGISTER_ERROR, "slaveDataSourceConfig配置未注入");
    }
    this.slaveDataSourceConfig = slaveDataSourceConfig;
    this.slaveSourceLoadBalancer = Optional.ofNullable(slaveSourceLoadBalancer)
      .orElse(ISlaveSourceLoadBalancer.of(WeightSlaveSourceLoadBalancer.class));
    // 设置目标数据源
    setTargetDataSources(slaves());
    // 设置默认数据源
    setDefaultTargetDataSource(master());
  }

  @Override
  protected Object determineCurrentLookupKey() {
    // 判断是读操作还是写操作
    try {
      if (DataSourceContextHolder.isRead()) {
        return slaveSourceLoadBalancer.load(slaveDataSourceConfig.getSlave());
      }
      return null;
    } catch (Exception e) {
      throw ServerException.build(ErrorEnum.SQL_EXECUTE_ERROR, e);
    } finally {
      DataSourceContextHolder.clear();
    }
  }

  private DataSource master() {
    DataSourceProperty master = slaveDataSourceConfig.getMaster();
    if (ObjectUtil.isEmpty(master)) {
      throw ServerException.build(ErrorEnum.BEAN_REGISTER_ERROR, "主库不能为空");
    }
    return master.newInstanceDataSource();
  }

  private Map<Object, Object> slaves() {
    return CollectionUtil.toMap(slaveDataSourceConfig.getSlave(), DataSourceProperty::getUnique,
      DataSourceProperty::newInstanceDataSource);
  }
}
